local function open_output(filename)
    local oldoutput = io.output()
    io.output(filename)
    local file = io.output()
    io.output(oldoutput)
    return file
end

local function write(message)
    OFP:write(message)
    OFP:flush()
end

function init() 

    --if INITIALIZED == nil then
        MUTATION_RATE_END = .025
        MUTATION_RATE_START = .005
        OFFSPRING_COUNT = 10
        MAX_GENERATION = 105

        --should not be greater than or equal to offspring count
        --or parents will be used for multiple generations
        PARENT_COUNT = 5
        GENDER_COUNT = 1 -- How many parents are combined to create child.
                         -- MUST BE LESS THAN OR EQUAL TO PARENT_COUNT!!

        INITIAL_GENE = "HOLD"  -- set to "RANDOM" for random parents
        SEED = 16 --get_world_state().frame_tick
        BOUNDARY = 3  -- How often to report approx average and stddev
        
        OFP = open_output("data/toribot_" .. GENDER_COUNT .. "_" .. SEED .. ".csv")
        write("GENERATION, GENAVG, ADJAVG, STDEV\n")

        -- players
        -- 0: red
        -- 1: blue
        PLAYER_NUM = 0
        GRIP_ENABLED = false

        --Bonus score added per additional dismember on opponent
        DISMEMBER_BONUS = 5000

        -- How sensitive the learner is to letting the opponent score
        OPPONENT_SCORE_MULTIPLIER = 0.5
        
        --if DISQUAL_ENABLED then
        --    run_cmd("loadmod toribot_disqual_mod.tbm") 
        --else
        --    run_cmd("loadmod toribot_mod.tbm")
        --end

        local rules = get_game_rules()
        run_cmd("chl 300")

        TURN_COUNT = 1 + math.floor((rules.matchframes-1) / rules.turnframes)
        -- echo("TURN_COUNT = "..TURN_COUNT)

        echo("seed: "..SEED)
        math.randomseed(SEED)
        CUR_OFFSPRING = 1
        GENERATION_COUNT = 1
        GENERATION_AVERAGE_COLLECTION = {}
        OFFSPRING = {}
        OFFSPRING_SCORE = {}
        seed_parents()
        for i=1,OFFSPRING_COUNT do
            OFFSPRING_SCORE[i]=0
        end

        MUTATION_RATE = function (turn) 
            return MUTATION_RATE_END * math.pow(0.8, -math.log(MUTATION_RATE_START/MUTATION_RATE_END)/math.log(0.8)/(TURN_COUNT-1) * (turn-TURN_COUNT)) 
        end

    
        INITIALIZED = true
--    end
end
--
-- Important info:
-- JOINT LIST
-- ==========
-- 0 = neck
-- 1 = chest
-- 2 = lumbar
-- 3 = abs
-- 4 = right pec
-- 5 = right shoulder
-- 6 = right elbow
-- 7 = left pec
-- 8 = left shoulder
-- 9 = left elbow
-- 10 = right wrist
-- 11 = left wrist
-- 12 = right glute
-- 13 = left glute
-- 14 = right hip
-- 15 = left hip
-- 16 = right knee
-- 17 = left knee
-- 18 = right ankle
-- 19 = left ankle
-- 
-- STATE LIST
-- ==========
-- 1 = extend/right rotate/right bend
-- 2 = contract/left rotate/left bend
-- 3 = hold
-- 4 = relax
-- 

function seed_parents()
    PARENTS = {}
    for parent=1, PARENT_COUNT do
        PARENTS[parent] = {gene={},value=0}

        for turn=1,TURN_COUNT do
            PARENTS[parent].gene[turn] ={}
            for joint=1,20 do
                if INITIAL_GENE == "RANDOM" then
                    PARENTS[parent].gene[turn][joint] = math.random(4) 
                else
                    PARENTS[parent].gene[turn][joint] = 3 
                end
            end

            if GRIP_ENABLED then
                for joint=21,22 do
                    if INITIAL_GENE == "RANDOM" then
                        PARENTS[parent].gene[turn][joint] = math.random(0,1)
                    else
                        PARENTS[parent].gene[turn][joint] = 0
                    end
                end
            end
        end
    end
    --echo("done making parents")
end

function cross_genes(genes, specific_parents)
    return genes[specific_parents[math.random(#specific_parents)]]
end

function cross_flex_genes(genes, turn, specific_parents)
    if math.random() < MUTATION_RATE(turn) then
        return math.random(4)
    else
        return cross_genes(genes, specific_parents)
    end
end

function cross_grip_genes(genes, turn, specific_parents)
    if math.random() < MUTATION_RATE(turn) then
        return math.random(0,1)
    else
        return cross_genes(genes, specific_parents)
    end
end

function generate_offspring()
    --echo("generation #"..GENERATION_COUNT)

    genes = {}
    for turn=1,TURN_COUNT do
        genes[turn] = {}
        for joint=1,20 do
            genes[turn][joint] = {}
            for parent=1,PARENT_COUNT do
                genes[turn][joint][parent]=PARENTS[parent].gene[turn][joint]
            end
        end
        
        if GRIP_ENABLED then

            for joint=21,22 do
                genes[turn][joint] = {}
                for parent=1,PARENT_COUNT do
                    genes[turn][joint][parent]=PARENTS[parent].gene[turn][joint]
                end
            end
        end
    end

    offspring = {}
    for i=1,OFFSPRING_COUNT do
        local specific_parents = {}
        local avail_parents = {}
        for parent=1,PARENT_COUNT do
            avail_parents[parent] = parent
        end
        for parent=1,GENDER_COUNT do
            specific_parents[parent] = 
                table.remove(avail_parents,math.random(#avail_parents))
        end

        offspring[i] = {}
        for turn=1,TURN_COUNT do
            offspring[i][turn] = {}
            for joint=1,20 do 
                offspring[i][turn][joint] = cross_flex_genes(genes[turn][joint], turn, specific_parents)
            end
            
            if GRIP_ENABLED then
                for joint=21,22 do 
                    offspring[i][turn][joint] = cross_grip_genes(genes[turn][joint], turn, specific_parents)
                end
            end

        end
    end

    return offspring
end

function on_new()
    CURRENT_TURN = 1
    --echo('Current Offspring: '..CUR_OFFSPRING)
    
    if CUR_OFFSPRING == 1 then
        OFFSPRING = generate_offspring()
        OFFSPRING_SCORE = {}
    end 
    on_freeze()
end

function on_freeze()
    --echo ("current turn:" .. CURRENT_TURN)
    for i=1,20 do
        set_joint_state(PLAYER_NUM, i-1, OFFSPRING[CUR_OFFSPRING][CURRENT_TURN][i]) -- JOINTS ARE 0 INDEXED!!
    end

    if GRIP_ENABLED then
        for i=1,2 do
            set_grip_info(PLAYER_NUM, i+10, OFFSPRING[CUR_OFFSPRING][CURRENT_TURN][i+20]) -- JOINTS ARE 0 INDEXED!!
        end
    end

    CURRENT_TURN = CURRENT_TURN + 1
end



function max_scores()
    
    --Find the best PARENT_COUNT genes among the OFFSPRING and PARENTS
    for i=1,OFFSPRING_COUNT do
        o=1
        --Find the first parent (o) to replace with i
        while o <= PARENT_COUNT do

            if OFFSPRING_SCORE[i] > PARENTS[o].value then
                p = PARENT_COUNT
                -- Shuffle all remaining parents down the list
                while p > o do
                    PARENTS[p] = PARENTS[p-1]
                    p = p-1
                end
                PARENTS[o] = {gene = OFFSPRING[i], value=OFFSPRING_SCORE[i]}
                o = PARENT_COUNT
            end
            o = o + 1
        end
        
    end

    return max
end

function on_end(mode)
    --echo("Winner via "..mode..":")
    -- 0: timeout
    -- 1: tie
    -- 2: disqual
	local myscore = get_player_info(1-PLAYER_NUM).injury
	local oppscore = get_player_info(PLAYER_NUM).injury

	
--	if myscore > oppscore then
--		echo ("I @ "..myscore.." beat OPPONENT @ "..oppscore)
--    elseif oppscore > myscore then
--		echo ("OPPONENT @ "..oppscore.." beat ME @ "..myscore)
--    else
--        echo ("OPPONENT @ "..oppscore.." tied ME @ "..myscore)
--	end
    
    -- EXPERIMENTAL: NEED SOME WAY TO DETERMINE WHO WINS A DQ
    --disqual_bonus = 0
    --if mode == 2 and then
    --    disqual_bonus = SELF_DISQUAL_PENALTY
    --elseif mode == 2 and then
    --    disqual_bonus = ENEMY_DISQUAL_BONUS
    --end

    dismember_count = 0
    for i=0,19 do
        if get_joint_dismember(1-PLAYER_NUM,i) then
            dismember_count = 1 + dismember_count 
        end
    end

    OFFSPRING_SCORE[CUR_OFFSPRING] = myscore 
    - oppscore * OPPONENT_SCORE_MULTIPLIER 
    + dismember_count * DISMEMBER_BONUS 

    --echo("I scored "..OFFSPRING_SCORE[CUR_OFFSPRING])

    CUR_OFFSPRING = CUR_OFFSPRING + 1
    if CUR_OFFSPRING > OFFSPRING_COUNT then
        CUR_OFFSPRING = 1
        sum = 0
        for i, score in pairs(OFFSPRING_SCORE) do
            sum = sum + score
        end
        
        if      (GENERATION_COUNT % BOUNDARY == 0 or
                (GENERATION_COUNT + 1) % BOUNDARY == 0 or
                (GENERATION_COUNT - 1) % BOUNDARY == 0 ) then
            GENERATION_AVERAGE_COLLECTION[#GENERATION_AVERAGE_COLLECTION + 1] = sum/OFFSPRING_COUNT
        end

        if GENERATION_COUNT > (BOUNDARY - 1) and (GENERATION_COUNT - 1) % BOUNDARY == 0 then
            local avgsum = 0
            local count = 0
            for k,v in pairs(GENERATION_AVERAGE_COLLECTION) do
                avgsum = v + avgsum
                count = count + 1
            end

            local avg = avgsum/count

            local d = 0

            for k,v in pairs(GENERATION_AVERAGE_COLLECTION) do
                d = v - avg
                avgsum = avgsum + (d * d)
            end

            std = math.sqrt(avgsum / count)

            echo("Approx average around "..(GENERATION_COUNT - 1).." is " .. avg.. " with stddev " .. std)
            write("," .. avg .. "," .. std .. "\n")
            GENERATION_AVERAGE_COLLECTION = {}
        elseif GENERATION_COUNT > 1 then
            write(",,\n")
        end

        echo("Generation #"..GENERATION_COUNT.." average score is "..sum/OFFSPRING_COUNT)
        write(GENERATION_COUNT .. "," .. (sum / OFFSPRING_COUNT))

        GENERATION_COUNT = GENERATION_COUNT + 1
        if GENERATION_COUNT > MAX_GENERATION then
            remove_hooks("toribot")
            OFP:write(",,\n")
            OFP:close()
        end

        max_scores()
        --echo("Mixing genes from individuals:")
        for i=1,OFFSPRING_COUNT do
            OFFSPRING_SCORE[i]=0
        end
    end
    -- reset game, calls on_new
    run_cmd("rt")

end

function debug()
    state = get_world_state()
    echo("Frame: " .. state.match_frame)
    echo("Turn: " .. CURRENT_TURN)
    echo("Generation: ".. GENERATION_COUNT)
    echo("Offspring: " .. CUR_OFFSPRING)
end

init()

add_hook("enter_freeze", "toribot", on_freeze)
add_hook("end_game", "toribot", on_end)
add_hook("new_game", "toribot", on_new)
--add_hook("enter_frame", "test", debug)
run_cmd("rt")